from utils.units import ConvUnit, ConvRelu, ConvPool, Dense
from torchsummary import summary


def test_conv_unit():
    conv_unit = ConvUnit(3, 64)
    print(conv_unit)
    summary(conv_unit, (3, 32, 32))


def test_conv_relu():
    conv_relu = ConvRelu(3, 64)
    print(conv_relu)
    summary(conv_relu, (3, 32, 32))


def test_conv_pool():
    conv_pool = ConvPool(3, 64)
    print(conv_pool)
    summary(conv_pool, (3, 32, 32))


def test_dense():
    dense = Dense(10, 2)
    print(dense)
    summary(dense, (10,))


if __name__ == '__main__':
    test_dense()
